from dataclasses import dataclass
from typing import Dict, List, Optional
import json
import os

@dataclass
class RembgConfig:
    # 模型配置
    model_name: str = "u2net"  # 默认使用通用模型
    providers: List[str] = None  # 运行设备
    
    # Alpha matting 参数
    alpha_matting: bool = True  # 是否启用 alpha matting
    alpha_matting_foreground_threshold: int = 240  # 前景阈值
    alpha_matting_background_threshold: int = 10   # 背景阈值
    alpha_matting_erode_size: int = 10            # 腐蚀大小
    
    # 后处理参数
    post_process_mask: bool = True  # 是否启用蒙版后处理
    
    def __post_init__(self):
        if self.providers is None:
            self.providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    
    @classmethod
    def get_model_descriptions(cls) -> Dict[str, str]:
        return {
            "u2net": "通用模型 - 效果最好，适合各种场景，但速度较慢",
            "u2netp": "轻量级模型 - 速度快，效果稍差",
            "u2net_human_seg": "人像模型 - 专门针对人物图片优化",
            "silueta": "轮廓模型 - 速度最快，主要提取轮廓"
        }
    
    @classmethod
    def get_presets(cls) -> Dict[str, Dict]:
        return {
            "通用": {
                "model_name": "u2net",
                "alpha_matting": True,
                "alpha_matting_foreground_threshold": 240,
                "alpha_matting_background_threshold": 10,
                "alpha_matting_erode_size": 10,
                "post_process_mask": True
            },
            "人像": {
                "model_name": "u2net_human_seg",
                "alpha_matting": True,
                "alpha_matting_foreground_threshold": 250,
                "alpha_matting_background_threshold": 5,
                "alpha_matting_erode_size": 15,
                "post_process_mask": True
            },
            "产品": {
                "model_name": "u2net",
                "alpha_matting": True,
                "alpha_matting_foreground_threshold": 245,
                "alpha_matting_background_threshold": 8,
                "alpha_matting_erode_size": 8,
                "post_process_mask": True
            },
            "快速": {
                "model_name": "u2netp",
                "alpha_matting": False,
                "post_process_mask": True
            }
        } 